#include <stdio.h>
#include <math.h>
#include "affine_warp.h"

void pseudoInverse(float* transform_matrix, float* result) {
    float a = transform_matrix[0];
    float b = transform_matrix[1];
    float c = transform_matrix[2];
    float d = transform_matrix[3];
    float e = transform_matrix[4];
    float f = transform_matrix[5];

    float determinant = a * e - b * d;

    // 计算伪逆矩阵
    if (fabs(determinant) > 0.0001) {
        result[0] =  e / determinant;
        result[1] = -b / determinant;
        result[3] = -d / determinant;
        result[4] =  a / determinant;
        result[2] = (c * e - b * f) / determinant;
        result[5] = (a * f - c * d) / determinant;
    } else {
        printf("Matrix is singular, cannot be inverted.\n");
    }
}


void affineWarpRGBLeftTopNeighbor(RGBPixel* input_image, RGBPixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;
    int new_index;

    // 对每个像素进行变换
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].r = input_image[new_index].r;
                output_image[index].g = input_image[new_index].g;
                output_image[index].b = input_image[new_index].b;
            }
        }
    }
}

void affineWarpYUV422CenterBili(YUV422Pixel* input_image, YUV422Pixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;

    // 计算图像中心点坐标
    float center_x = width / 2.0f;
    float center_y = height / 2.0f;

    transform_matrix[2] = transform_matrix[0] * (-center_x) + transform_matrix[1] * (-center_y) + center_x;
    transform_matrix[5] = transform_matrix[3] * (-center_x) + transform_matrix[4] * (-center_y) + center_y;

    // 对每个像素进行变换
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].y = input_image[new_index].y;

                output_image[index].uv = input_image[new_index].uv;
            }
        }
    }
}

void affineWarpYUV422LeftTopBili(YUV422Pixel* input_image, YUV422Pixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;

    // 对每个像素进行变换
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].y = (1 - dy) * ((1 - dx) * input_image[new_index].y
                                                    + dx * input_image[new_index + 1].y)
                                        + dy * ((1 - dx) * input_image[new_index + width].y
                                                + dx * input_image[new_index + width + 1].y);

                output_image[index].uv = (1 - dy) * ((1 - dx) * input_image[new_index].uv
                                                     + dx * input_image[new_index + 2].uv)
                                         + dy * ((1 - dx) * input_image[new_index + width].uv
                                                 + dx * input_image[new_index + width + 2].uv);
            }
        }
    }
}

void affineWarpYUV422CenterBiliComplex(YUV422Pixel* input_image, YUV422Pixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;

    // 计算图像中心点坐标
    float center_x = width / 2.0f;
    float center_y = height / 2.0f;

    transform_matrix[2] = transform_matrix[0] * (-center_x) + transform_matrix[1] * (-center_y) + center_x;
    transform_matrix[5] = transform_matrix[3] * (-center_x) + transform_matrix[4] * (-center_y) + center_y;

    // 对每个像素进行变换
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].y = (1 - dy) * ((1 - dx) * input_image[new_index].y + dx * input_image[new_index + 1].y)
                                        + dy * ((1 - dx) * input_image[new_index + width].y + dx * input_image[new_index + width + 1].y);

                if(new_index % 2 == 0) {
                    //u
                    output_image[index].uv = (1 - dy) * ((1 - dx) * input_image[new_index].uv + dx * input_image[new_index + 2].uv)
                                             + dy * ((1 - dx) * input_image[new_index + width].uv + dx * input_image[new_index + width + 2].uv);
                } else {
                    //v
                    output_image[index].uv = (1 - dy) * ((1 - dx) * input_image[new_index].uv + dx * input_image[new_index + 2].uv)
                                             + dy * ((1 - dx) * input_image[new_index + width].uv + dx * input_image[new_index + width + 2].uv);
                }


//                output_image[index].y = input_image[new_index].y;
//                output_image[index].uv = input_image[new_index].uv;
            }
        }
    }
}
void affineWarpYUV422CenterNeibor(YUV422Pixel* input_image, YUV422Pixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;

    // 计算图像中心点坐标
    float center_x = width / 2.0f;
    float center_y = height / 2.0f;

    transform_matrix[2] = transform_matrix[0] * (-center_x) + transform_matrix[1] * (-center_y) + center_x;
    transform_matrix[5] = transform_matrix[3] * (-center_x) + transform_matrix[4] * (-center_y) + center_y;

    // 对每个像素进行变换
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                output_image[index].y = input_image[new_index].y;
                output_image[index].uv = input_image[new_index].uv;
            }
        }
    }
}
void affineWarpRGBLeftTopBili(RGBPixel* input_image, RGBPixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;
    int new_index;

    // 对每个像素进行变换
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].r = (1 - dy) * ((1 - dx) * input_image[new_index].r + dx * input_image[new_index + 1].r)
                                        + dy * ((1 - dx) * input_image[new_index + width].r + dx * input_image[new_index + width + 1].r);

                output_image[index].g = (1 - dy) * ((1 - dx) * input_image[new_index].g + dx * input_image[new_index + 1].g)
                                        + dy * ((1 - dx) * input_image[new_index + width].g + dx * input_image[new_index + width + 1].g);

                output_image[index].b = (1 - dy) * ((1 - dx) * input_image[new_index].b + dx * input_image[new_index + 1].b)
                                        + dy * ((1 - dx) * input_image[new_index + width].b + dx * input_image[new_index + width + 1].b);
            }
        }
    }
}

void affineWarpYUV444LeftTopBili(YUV444Pixel* input_image, YUV444Pixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;
    int new_index;

    // 对每个像素进行变换
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].y = (1 - dy) * ((1 - dx) * input_image[new_index].y + dx * input_image[new_index + 1].y)
                                        + dy * ((1 - dx) * input_image[new_index + width].y + dx * input_image[new_index + width + 1].y);

                output_image[index].u = (1 - dy) * ((1 - dx) * input_image[new_index].u + dx * input_image[new_index + 1].u)
                                        + dy * ((1 - dx) * input_image[new_index + width].u + dx * input_image[new_index + width + 1].u);

                output_image[index].v = (1 - dy) * ((1 - dx) * input_image[new_index].v + dx * input_image[new_index + 1].v)
                                        + dy * ((1 - dx) * input_image[new_index + width].v + dx * input_image[new_index + width + 1].v);
            }
        }
    }
}

void affineWarpRGBCenterBili(RGBPixel* input_image, RGBPixel* output_image, int width, int height, float* transform_matrix) {
    float x, y, new_x, new_y;

    // 组合变换矩阵
    // 先将图像中心平移到原点，旋转后再移回原来位置 TransforMatrix =  T(p) R(α) T(-p);

    // 计算图像中心点坐标
    float center_x = width / 2.0f;
    float center_y = height / 2.0f;

//    // 将图像中心设置为旋转中心α
//    float translate_to_origin[6] = {
//            1.0, 0.0, -center_x,
//            0.0, 1.0, -center_y
//    };
//    // 将图片再平移回去
//    float translate_back[6] = {
//            1.0, 0.0, center_x,
//            0.0, 1.0, center_y
//    };

    transform_matrix[2] = transform_matrix[0] * (-center_x) + transform_matrix[1] * (-center_y) + center_x;
    transform_matrix[5] = transform_matrix[3] * (-center_x) + transform_matrix[4] * (-center_y) + center_y;

    // 对每个像素进行变换
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            x = j;
            y = i;

            // 计算变换后的坐标
            new_x = transform_matrix[0] * x + transform_matrix[1] * y + transform_matrix[2];
            new_y = transform_matrix[3] * x + transform_matrix[4] * y + transform_matrix[5];

            // 获取整数部分坐标
            int int_x = (int)new_x;
            int int_y = (int)new_y;

            // 双线性插值
            float dx = new_x - int_x;
            float dy = new_y - int_y;

            if (int_x >= 0 && int_x < width - 1 && int_y >= 0 && int_y < height - 1) {
                int index = i * width + j;
                int new_index = (int_y * width) + int_x;

                // 计算新像素值
                output_image[index].r = (1 - dy) * ((1 - dx) * input_image[new_index].r + dx * input_image[new_index + 1].r)
                                        + dy * ((1 - dx) * input_image[new_index + width].r + dx * input_image[new_index + width + 1].r);

                output_image[index].g = (1 - dy) * ((1 - dx) * input_image[new_index].g + dx * input_image[new_index + 1].g)
                                        + dy * ((1 - dx) * input_image[new_index + width].g + dx * input_image[new_index + width + 1].g);

                output_image[index].b = (1 - dy) * ((1 - dx) * input_image[new_index].b + dx * input_image[new_index + 1].b)
                                        + dy * ((1 - dx) * input_image[new_index + width].b + dx * input_image[new_index + width + 1].b);
            }
        }
    }
}




/****************************************************************************************/
/******************************* 调用opencv进行图像显示 ************************************/
/****************************************************************************************/

#if 0
void testSrcTgtImg(const char* file_path) {
    // 读取RGB图片
    cv::Mat image = cv::imread(file_path, cv::IMREAD_COLOR);

    if (image.empty()) {
        printf("Could not open or find the image.\n");
        return;
    }

    int width = image.cols;
    int height = image.rows;

    // 创建另一张图片，用于接收R、G、B值
    cv::Mat output_image(height, width, CV_8UC3);

    // 分别赋值R、G、B值给接收图片
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            cv::Vec3b pixel = image.at<cv::Vec3b>(y, x);

            // 分别赋值给另一张图片的R、G、B通道
            output_image.at<cv::Vec3b>(y, x)[0] = pixel[0]; // B值
            output_image.at<cv::Vec3b>(y, x)[1] = pixel[1]; // G值
            output_image.at<cv::Vec3b>(y, x)[2] = pixel[2]; // R值
        }
    }

    // 显示原始图片和接收图片
    cv::imshow("Original Image", image);
    cv::imshow("Output Image (R/G/B channels)", output_image);
    cv::waitKey(0);
    cv::destroyAllWindows();

}

void testAffineWarpRGB(const char* picPath) {

    float transform_matrix[6] = {
            0.866, -0.5, 0.0,
            0.5, 0.866, 0.0
    };

    // 图像宽高
    int width = 800;
    int height = 600;

    // 分配内存存储输入和输出图像
    RGBPixel* input_image = (RGBPixel*)malloc(width * height * sizeof(RGBPixel));
    RGBPixel* output_image = (RGBPixel*)malloc(width * height * sizeof(RGBPixel));

    cv::Mat image = cv::imread(picPath, cv::IMREAD_COLOR);

    if (image.empty()) {
        printf("Could not open or find the image.\n");
        return;
    }

    // 分别赋值R、G、B值给接收图片
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            cv::Vec3b pixel = image.at<cv::Vec3b>(y, x);
            int i = y * width + x;
            input_image[i].r = pixel[2];
            input_image[i].g = pixel[1];
            input_image[i].b = pixel[0];
        }
    }

//    affineWarpRGBLeftTopNeighbor(input_image, output_image, width, height, transform_matrix);
//    affineWarpRGBLeftTopBili(input_image, output_image, width, height, transform_matrix);
    affineWarpRGBCenterBili(input_image, output_image, width, height, transform_matrix);

    cv::Mat transformed_image(height, width, CV_8UC3);
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            transformed_image.at<cv::Vec3b>(i, j)[0] = output_image[i * width + j].b;
            transformed_image.at<cv::Vec3b>(i, j)[1] = output_image[i * width + j].g;
            transformed_image.at<cv::Vec3b>(i, j)[2] = output_image[i * width + j].r;
        }
    }

    // 使用 OpenCV 显示变换后的图像
    cv::imshow("Transformed Image", transformed_image);
    cv::waitKey(0);
    cv::destroyAllWindows();

    free(input_image);
    free(output_image);
}

void testAffineWarpYUV444(const char* file_path) {
    float transform_matrix[6] = {
            0.866, -0.5, 0.0,
            0.5, 0.866, 0.0
    };
//
//    float result[6] = {
//            0.0, 0.0, 0.0,
//            0.0, 0.0, 0.0
//    };
//    float transform_matrix[6] = {
//            1, 0, 0.0,
//            0, 1, 0.0
//    };

    // 图像宽高
    int width = 800;
    int height = 600;

    // 分配内存存储输入和输出图像
    YUV444Pixel* input_image_4 = (YUV444Pixel*)malloc(width * height * sizeof(YUV444Pixel));
    YUV444Pixel* output_image_4 = (YUV444Pixel*)malloc(width * height * sizeof(YUV444Pixel));

    FILE *file;
    char line[100]; // 假设每行最大长度为100
    int lineNumber = 0;

    file = fopen("demo422.txt", "r");
    if (file == NULL) {
        perror("Error opening the file");
        return;
    }

    while (fgets(line, sizeof(line), file) != NULL) {
        unsigned int value;
        if (sscanf(line, "%x", &value) != 1) {
            printf("Error reading line %d\n", lineNumber);
            continue;
        }

        if (lineNumber % 2 == 0) {
            input_image_4[lineNumber / 2].y = (unsigned char)value;
        } else {
            if((lineNumber / 2) % 2 == 0) {
                //u
                input_image_4[lineNumber / 2].u = (unsigned char)value;
            }else{
                //v
                input_image_4[lineNumber / 2].v = (unsigned char)value;
            }
        }
        lineNumber++;
    }

    for(int k = 0; k < width * height - 1; k++) {
        if(k % 2 == 0) {
            input_image_4[k].v = input_image_4[k + 1].v;
            input_image_4[k + 1].u = input_image_4[k].u;
        }
    }

    fclose(file);
    printf("input_image lineNumber: %d\n",lineNumber);


    affineWarpYUV444LeftTopBili(input_image_4, output_image_4, width, height, transform_matrix);


    // 调用opencv显示一下图片,需先拼装opencv rgb，再展示
    cv::Mat transformed_image(height, width, CV_8UC3);
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {

            unsigned char y, u, v;
            y = output_image_4[i * width + j].y;
            u = output_image_4[i * width + j].u;
            v = output_image_4[i * width + j].v;

            //yu组
            //b
            transformed_image.at<cv::Vec3b>(i, j)[0] = (uint8_t)(y + (u - 128) + ((104 * (u - 128)) >> 8));
            //g
            transformed_image.at<cv::Vec3b>(i, j)[1] = (uint8_t)(y - (89 * (v - 128) >> 8) - ((183 * (u - 128)) >> 8));
            //r
            transformed_image.at<cv::Vec3b>(i, j)[2] = (uint8_t)(y + (v - 128) + ((199 * (v - 128)) >> 8));

        }
    }

    cv::imshow("Transformed Image", transformed_image);
    cv::waitKey(0);
    cv::destroyAllWindows();

    free(input_image_4);
    free(output_image_4);

}

void testAffineWarpYUV422(const char* filePath) {

//    float transform_matrix[6] = {
//            0.866, -0.5, 0.0,
//            0.5, 0.866, 0.0
//    };

//    float transform_matrix[6] = {
//            0.5, -0.866, 0.0,
//            0.866, 0.5, 0.0
//    };

//    float result[6] = {
//            0.0, 0.0, 0.0,
//            0.0, 0.0, 0.0
//    };
    float transform_matrix[6] = {
            1, 0, 0.0,
            0, 1, 0.0
    };
//
//    float transform_matrix[6] = {
//            0.996, -0.09, 0.0,
//            0.09, 0.996, 0.0
//    };

//    float transform_matrix[6] = {
//            0.998, -0.052, 0.0,
//            0.052, 0.998, 0.0
//    };

    // 图像宽高
    int width = 800;
    int height = 600;

    // 分配内存存储输入和输出图像
    YUV422Pixel* input_image = (YUV422Pixel*)malloc(width * height * sizeof(YUV422Pixel));
    YUV422Pixel* output_image = (YUV422Pixel*)malloc(width * height * sizeof(YUV422Pixel));

    FILE *file;
    char line[100]; // 假设每行最大长度为100
    int lineNumber = 0;

    file = fopen(filePath, "r");
    if (file == NULL) {
        perror("Error opening the file");
        return;
    }

    while (fgets(line, sizeof(line), file) != NULL) {
        unsigned int value;
        if (sscanf(line, "%x", &value) != 1) {
            printf("Error reading line %d\n", lineNumber);
            continue;
        }

        if (lineNumber % 2 == 0) {
            input_image[lineNumber / 2].y = (unsigned char)value;
        } else {
            input_image[lineNumber / 2].uv = (unsigned char)value;
        }

        lineNumber++;
    }

    fclose(file);
    printf("input_image lineNumber: %d\n",lineNumber);

    affineWarpYUV422LeftTopBili(input_image, output_image, width, height, transform_matrix);

//  变换矩阵取逆
//    pseudoInverse(transform_matrix, result);
//    affineWarpYUV422LeftTopBiliComlex(input_image, output_image, width, height, result);

//  变换矩阵不取逆
//    affineWarpYUV422CenterBiliComplex(input_image, output_image, width, height, transform_matrix);

    // 调用opencv显示一下图片,需先拼装opencv rgb，再展示
    cv::Mat transformed_image(height, width, CV_8UC3);
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            if(j % 2 == 1) continue;

            unsigned char y1, y2, u, v;
            y1 = output_image[i * width + j].y;
            u = output_image[i * width + j].uv;
            y2 = output_image[i * width + j + 1].y;
            v = output_image[i * width + j + 1].uv;

//            y1 = input_image[i * width + j].y;
//            u = input_image[i * width + j].uv;
//            y2 = input_image[i * width + j + 1].y;
//            v = input_image[i * width + j + 1].uv;

            //yu组
            //b
            transformed_image.at<cv::Vec3b>(i, j)[0] = (uint8_t)(y1 + (u - 128) + ((104 * (u - 128)) >> 8));
            //g
            transformed_image.at<cv::Vec3b>(i, j)[1] = (uint8_t)(y1 - (89 * (v - 128) >> 8) - ((183 * (u - 128)) >> 8));
            //r
            transformed_image.at<cv::Vec3b>(i, j)[2] = (uint8_t)(y1 + (v - 128) + ((199 * (v - 128)) >> 8));

            //yv组
            //b
            transformed_image.at<cv::Vec3b>(i, j + 1)[0] = (uint8_t)(y2 + (u - 128) + ((104 * (u - 128)) >> 8));
            //g
            transformed_image.at<cv::Vec3b>(i, j + 1)[1] = (uint8_t)(y2 - (89 * (v - 128) >> 8) - ((183 * (u - 128)) >> 8));
            //r
            transformed_image.at<cv::Vec3b>(i, j + 1)[2] = (uint8_t)(y2 + (v - 128) + ((199 * (v - 128)) >> 8));
        }
    }

    cv::imshow("Transformed Image", transformed_image);
    cv::waitKey(0);
    cv::destroyAllWindows();

    free(input_image);
    free(output_image);
}
#endif

